import numpy as np
import cv2
import glob
import json

def compute_extrinsics(A, homographies):
    # 计算所有图像的平移和旋转向量
    num_images = len(homographies)
    rotation_vectors = np.zeros((3, num_images))
    translation_vectors = np.zeros((3, num_images))
    A_inv = np.linalg.inv(A)

    for i in range(num_images):
        H = homographies[i]
        h1 = H[:, 0]
        h2 = H[:, 1]
        h3 = H[:, 2]
        lambda_ = 1 / np.linalg.norm(A_inv @ h1)

        # 3D 旋转矩阵
        r1 = lambda_ * A_inv @ h1
        r2 = lambda_ * A_inv @ h2
        r3 = np.cross(r1, r2)
        R = np.column_stack((r1, r2, r3))

        # 使用OpenCV将旋转矩阵转换为旋转向量
        rotation_vector, _ = cv2.Rodrigues(R)
        rotation_vectors[:, i] = rotation_vector.flatten()

        # 平移向量
        t = lambda_ * A_inv @ h3
        translation_vectors[:, i] = t

    rotation_vectors = rotation_vectors.T
    translation_vectors = translation_vectors.T

    return rotation_vectors, translation_vectors


def parse_inputs(K, P):
    # 解析输入参数并初始化相机模型和校准参数

    world_units = "mm"  # 世界单位，通常为毫米
    camera_model = {}
    calibration_params = {}

    # 初始化相机模型参数
    camera_model = {
        "estimate_skew": False,  # 是否估计倾斜
        "estimate_tangential_distortion": False,  # 是否估计切向畸变
        "num_radial_distortion_coefficients": 2,  # 径向畸变系数的数量
    }

    # 初始化内参矩阵
    init_intrinsics = np.array(K, dtype=np.double)

    # 初始化径向畸变系数
    init_radial = np.array(P, dtype=np.double)

    # 初始化校准参数
    calibration_params["init_intrinsics"] = init_intrinsics
    calibration_params["init_radial"] = init_radial
    calibration_params["show_progress_bar"] = 0  # 是否显示进度条

    return world_units, camera_model, calibration_params


def compute_homographies(image_points, world_points):
    # 计算所有图像的单应性矩阵

    # 由于Python没有与MATLAB相同的警告处理机制，我们不需要模拟警告
    num_images = len(image_points)
    valid_idx = np.full(num_images, True, dtype=bool)
    homographies = [[] for _ in range(1)]
    for i in range(num_images):
        try:
            # 在Python中，我们直接使用numpy数组，不需要将它们转换为double
            homographies[i], mask = cv2.findHomography(
                world_points, image_points[i], cv2.RANSAC, 5.0
            )
        except:
            valid_idx[i] = False

    # 只保留有效的单应性矩阵
    homographies = [h for i, h in enumerate(homographies) if valid_idx[i]]
    num_invalid = np.sum(~valid_idx)
    if num_invalid > 0:
        print(
            f"警告：{num_invalid}个单应性矩阵无效，共计算了{num_images}个单应性矩阵。"
        )

    return homographies, valid_idx


def compute_initial_parameter_estimate(
    image_points, world_points, camera_model, world_units, init_intrinsics, init_radial
):
    # 计算单应性矩阵
    H, valid_idx = compute_homographies(image_points, world_points)

    # 初始化内参矩阵
    A = init_intrinsics

    # 计算外参（旋转和平移向量）
    rvecs, tvecs = compute_extrinsics(A, H)

    # 初始化径向畸变系数
    radial_coeffs = init_radial

    # 创建相机参数结构
    initial_params = {
        "intrinsic_matrix": A,
        "rotation_vectors": rvecs,
        "translation_vectors": tvecs,
        "world_points": world_points,
        "world_units": world_units,
        "estimate_skew": camera_model["estimate_skew"],
        "num_radial_distortion_coefficients": camera_model[
            "num_radial_distortion_coefficients"
        ],
        "estimate_tangential_distortion": camera_model[
            "estimate_tangential_distortion"
        ],
        "radial_distortion": radial_coeffs,
    }

    return initial_params, valid_idx


def calibrate_one_camera(
    image_points, world_points, camera_model, world_units, calibration_params
):
    # print("开始计算初始参数估计...")

    # 计算初始参数估计
    camera_params, images_used = compute_initial_parameter_estimate(
        image_points,
        world_points,
        camera_model,
        world_units,
        calibration_params["init_intrinsics"],
        calibration_params["init_radial"],
    )
    # imagePoints = imagePoints[:, :, imagesUsed];

    # print("开始优化相机参数和计算畸变系数...")

    return camera_params, images_used


def estimate_camera_parameters(K, P, image_points, world_points):
    # 解析输入参数并初始化相机模型和校准参数
    world_units, camera_model, calibration_params = parse_inputs(K, P)

    # 校准相机并获取校准参数和使用的图像
    camera_params, images_used = calibrate_one_camera(
        image_points, world_points, camera_model, world_units, calibration_params
    )

    return camera_params, images_used


# 设置寻找亚像素角点的参数，采用的停止准则是最大循环次数30和最大误差容限0.001
criteria = (cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 30, 0.001)
# 设置标定板的大小和角点数量
m = 11
n = 8
squaresize = 20
# 创建一个只包含XY坐标的数组
objp = np.zeros((m * n, 2), np.float32)
# 生成XY网格
for i in range(n):
    for j in range(m):
        # 每个点的X和Y坐标（横向排列）
        objp[i * m + j, :2] = (j, i)
objp = objp * squaresize

K = np.array(
    [
        [1208.60235584270, 0, 733.966799900930],
        [0, 1208.47260369893, 554.618880013422],
        [0, 0, 1],
    ],
    dtype=np.float64,
)
P = np.array([-0.0554732991865857, 0.0826537426163534], dtype=np.float64)
obj_points = []  # 存储3D点
img_points = []  # 存储2D点


def UpdateObjp():
    objp = np.zeros((m * n, 2), np.float32)
    # 生成XY网格
    for i in range(n):
        for j in range(m):
            # 每个点的X和Y坐标（横向排列）
            objp[i * m + j, :2] = (j, i)
    objp = objp * squaresize


def SetBoardSizeM(val: int):
    m = val


def GetBoardSizeM() -> int:
    return m


def SetBoardSizeN(val: int):
    n = val


def GetBoardSizeN() -> int:
    return n


def SetSquareSize(val: int):
    squaresize = val


def GetSquaresize() -> int:
    return squaresize


def HandleImage(gray):
    UpdateObjp()
    size = gray.shape[::-1]
    ret, corners = cv2.findChessboardCorners(gray, (11, 8), None)
    if ret:
        obj_points.append(objp)
        corners2 = cv2.cornerSubPix(
            gray, corners, (5, 5), (-1, -1), criteria
        )  # 在原角点的基础上寻找亚像素角点
        if [corners2]:
            # img_points.append(corners2)
            img_points = [corners2]
        else:
            # img_points.append(corners)
            img_points = [corners]
    for i in range(len(img_points)):
        img_points[i] = img_points[i].reshape((-1, 2))
    camera_params1, images_used = estimate_camera_parameters(
        K, P, img_points, objp)
    B = camera_params1["translation_vectors"]
    jl = []
    for i in range(len(B)):
        # distance = np.linalg.norm(B[i] -B[0])
        distance = np.linalg.norm(B[i])
        jl.append(distance)
    print(jl)
    return str(jl[0])
